from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, SkipExhausted


class SlidingWindowIterator(DatasetIterator[tuple[int, int]]):
    def __init__(self, *,
        start: int = 0,
        stop: int | None = None,
        window: int,
        stride: int,
        remainder: bool = False,
    ):
        self._start = start
        self._stop = stop
        self._window = window
        self._stride = stride
        self._remainder = remainder

    @property
    def start(self) -> int:
        return self._start

    @property
    def stop(self) -> int | None:
        return self._stop

    @property
    def window(self) -> int:
        return self._window

    @property
    def stride(self) -> int:
        return self._stride

    @property
    def remainder(self) -> bool:
        return self._remainder

    def __next__(self) -> tuple[int, int]:
        if self._stop is not None and self._stop - self._start < self._window:
            if self._remainder and self._stop - self._start > 0:
                head = self._start
                tail = self._stop
                self._start = self._stop
                return head, tail
            raise StopIteration
        head = self._start
        tail = head + self._window
        self._start += self._stride
        return head, tail

    def __skip__(self, skip_windows_n: int):
        if self._stop is not None:
            windows_n = _compute_windows_n(
                start=self._start,
                stop=self._stop,
                window=self._window,
                stride=self._stride,
                remainder=self.remainder)
            if skip_windows_n > windows_n:
                self._start += windows_n * self._stride
                raise SkipExhausted(windows_n)
        self._start += skip_windows_n * self._stride

    def __pause__(self) -> int:
        return self._start


class SlidingWindowDataset(Dataset[tuple[int, int]]):
    def __init__(self, *,
        start: int = 0,
        stop: int | None = None,
        window: int,
        stride: int,
        remainder: bool = False,
    ):
        if window < 0:
            raise ValueError(f"Expected window>=0 got window={window}")
        if stride <= 0:
            raise ValueError(f"Expected stride>0 got stride={stride}")

        self._start = start
        self._stop = stop
        self._window = window
        self._stride = stride
        self._remainder = remainder

    @property
    def start(self) -> int:
        return self._start

    @property
    def stop(self) -> int | None:
        return self._stop

    @property
    def window(self) -> int:
        return self._window

    @property
    def stride(self) -> int:
        return self._stride

    @property
    def remainder(self) -> bool:
        return self._remainder

    def __len__(self) -> int:
        if self._stop is None:
            raise TypeError
        return _compute_windows_n(
            start=self.start,
            stop=self.stop,
            window=self.window,
            stride=self.stride,
            remainder=self.remainder)

    def __getitem__(self, index: int, /) -> tuple[int, int]:
        if index < 0:
            raise IndexError(f'Expected 0<=index, got {index}.')
        if self._stop is not None and index >= len(self):
            raise IndexError(f'Expected index<{len(self)}, got {index}.')

        head = index * self._stride + self._start
        tail = head + self._window
        return head, tail

    def __iter__(self) -> SlidingWindowIterator:
        return SlidingWindowIterator(
            start=self._start,
            stop=self._stop,
            window=self._window,
            stride=self._stride,
            remainder=self._remainder)

    def __resume__(self, state: int | None = None, /) -> SlidingWindowIterator:
        return SlidingWindowIterator(
            start=self._start if state is None else state,
            stop=self._stop,
            window=self._window,
            stride=self._stride,
            remainder=self._remainder)


def _compute_windows_n(*,
    start: int = 0,
    stop: int,
    window: int,
    stride: int,
    remainder: bool = False,
) -> int:
    items_n = stop - start
    if items_n < window:
        if remainder and items_n > 0:
            return 1
        return 0

    items_n -= window
    segments_n = 1 + items_n // stride
    if remainder and items_n % stride > 0:
        segments_n += 1
    return segments_n
